#!/usr/bin/env python3
"""
安全帽检测模型训练脚本
使用YOLOv8训练安全帽检测模型
"""

import os
import yaml
from ultralytics import YOLO
import torch
import matplotlib.pyplot as plt
import seaborn as sns
from pathlib import Path

def setup_environment():
    """设置训练环境"""
    # 检查CUDA是否可用
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"使用设备: {device}")
    
    # 检查GPU信息
    if torch.cuda.is_available():
        print(f"GPU名称: {torch.cuda.get_device_name(0)}")
        print(f"GPU内存: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
    
    return device

def check_dataset():
    """检查数据集配置"""
    data_config = "data.yaml"
    
    if not os.path.exists(data_config):
        print(f"错误: 找不到数据集配置文件 {data_config}")
        return False
    
    with open(data_config, 'r', encoding='utf-8') as f:
        data = yaml.safe_load(f)
    
    print("数据集信息:")
    print(f"类别数量: {data['nc']}")
    print(f"类别名称: {data['names']}")
    
    # 检查路径
    train_path = data['train'].replace('../', '')
    val_path = data['val'].replace('../', '')
    test_path = data['test'].replace('../', '')
    
    print(f"训练集路径: {train_path}")
    print(f"验证集路径: {val_path}")
    print(f"测试集路径: {test_path}")
    
    return True

def train_model():
    """训练YOLO模型"""
    print("开始训练YOLO8安全帽检测模型...")
    
    # 加载预训练模型
    model = YOLO('yolov8n.pt')  # 使用nano版本，速度快
    
    # 训练配置
    training_config = {
        'data': 'data.yaml',
        'epochs': 100,
        'imgsz': 640,
        'batch': 16,
        'device': 'auto',
        'workers': 4,
        'patience': 15,
        'save': True,
        'save_period': 10,
        'cache': False,
        'optimizer': 'auto',
        'lr0': 0.01,
        'lrf': 0.01,
        'momentum': 0.937,
        'weight_decay': 0.0005,
        'warmup_epochs': 3,
        'warmup_momentum': 0.8,
        'warmup_bias_lr': 0.1,
        'augment': True,
        'mixup': 0.0,
        'copy_paste': 0.0,
        'project': 'runs/detect',
        'name': 'helmet_detection',
        'exist_ok': True,
        'verbose': True
    }
    
    # 开始训练
    results = model.train(**training_config)
    
    print("训练完成!")
    return results

def validate_model():
    """验证训练好的模型"""
    print("验证模型性能...")
    
    # 加载训练好的模型
    model_path = "runs/detect/helmet_detection/weights/best.pt"
    
    if not os.path.exists(model_path):
        print(f"错误: 找不到训练好的模型 {model_path}")
        return None
    
    model = YOLO(model_path)
    
    # 在验证集上评估
    metrics = model.val(data='data.yaml')
    
    print(f"mAP50: {metrics.box.map50:.3f}")
    print(f"mAP50-95: {metrics.box.map:.3f}")
    
    return metrics

def plot_training_results():
    """绘制训练结果"""
    results_dir = "runs/detect/helmet_detection"
    
    if not os.path.exists(results_dir):
        print("未找到训练结果目录")
        return
    
    # 设置绘图样式
    plt.style.use('seaborn-v0_8')
    sns.set_palette("husl")
    
    # 显示训练结果图片
    results_images = [
        "results.png",
        "confusion_matrix.png", 
        "train_batch0.jpg",
        "val_batch0_pred.jpg"
    ]
    
    fig, axes = plt.subplots(2, 2, figsize=(15, 10))
    axes = axes.flatten()
    
    for i, img_name in enumerate(results_images):
        img_path = os.path.join(results_dir, img_name)
        if os.path.exists(img_path):
            img = plt.imread(img_path)
            axes[i].imshow(img)
            axes[i].set_title(img_name.replace('.png', '').replace('.jpg', '').replace('_', ' ').title())
            axes[i].axis('off')
        else:
            axes[i].text(0.5, 0.5, f'{img_name}\n未找到', 
                        ha='center', va='center', transform=axes[i].transAxes)
            axes[i].axis('off')
    
    plt.tight_layout()
    plt.savefig('training_summary.png', dpi=300, bbox_inches='tight')
    plt.show()

def main():
    """主函数"""
    print("=" * 60)
    print("YOLO8 安全帽检测模型训练")
    print("=" * 60)
    
    # 设置环境
    device = setup_environment()
    
    # 检查数据集
    if not check_dataset():
        return
    
    # 训练模型
    try:
        results = train_model()
        
        # 验证模型
        metrics = validate_model()
        
        # 绘制结果
        plot_training_results()
        
        print("\n" + "=" * 60)
        print("训练完成! 最佳模型保存在: runs/detect/helmet_detection/weights/best.pt")
        print("=" * 60)
        
    except Exception as e:
        print(f"训练过程中出现错误: {e}")

if __name__ == "__main__":
    main() 